import protocol
from protocol import Protocol
import os
import time
import math

PACK_SIZE = 512

UPDATE_STA_IDLE = 0
UPDATE_STA_WAIT = 1
UPDATE_STA_UPLOAD = 2
UPDATE_STA_END = 3

UPDATE_EVENT = 0xa9
UPDATE_OP_REQ = 0
UPDATE_OP_DATA = 1
UPDATE_OP_END = 2

MODULE_MAX_SIZE = 128 * 1024


class MasterContorlUpdate():
    def __init__(self, msg):
        super().__init__()
        self.is_exit = False
        self.msg = msg
        self.protocol = Protocol()
        self.init_env()

    def init_env(self):
        self.app_path = ''
        self.app_pack = []
        self.err_count = 0
        self.start_time = time.time()

    def wait_mc_data(self, timeout_ms):
        end_time = time.time() + (timeout_ms / 1000)
        while end_time > time.time():
            ch = self.msg.read(1)
            if ch == None:
                continue
            if self.protocol.unpack(ch) == protocol.RET_SUCCESS:
                return True
        return False

    def get_mc_data(self):
        ''' 需要在解包成功后调用
            返回: 去掉包头的数据
        '''
        data = self.protocol.get_data()
        return data

    def load_updata_file(self, file_name):
        with open("/home/pi/sm2-fact-qc-mc/bin/"+file_name, "rb") as f:
            data = f.read()
            data = bytearray(data)
            pack_count = math.ceil(len(data) / PACK_SIZE)
            loop = len(data) // PACK_SIZE
            for i in range(loop):
                pack = bytearray([(i >> 8) & 0xff, i & 0xff])
                pack += data[i * PACK_SIZE : i * PACK_SIZE + PACK_SIZE]
                self.app_pack.append(pack)
            if loop != pack_count:
                pack = bytearray([(loop >> 8) & 0xff, loop & 0xff])
                pack += data[loop * PACK_SIZE : len(data)]
                self.app_pack.append(pack)
            self.pack_count = len(self.app_pack)

    def _failed_update(self):
        self.init_env()

    def _wait_update(self):
        if not self.msg._opened:
            time.sleep(1)
            return False
        self.msg.write(b"M2000 F100\r\n")
        self.msg.write(b"M1999\r\n")
        timeout = time.time() + 20
        while timeout > time.time():
            print("Wait Erasing flash")
            cmd = Protocol.pack(UPDATE_EVENT, UPDATE_OP_REQ)
            self.msg.write(cmd)
            if self.wait_mc_data(2000):
                data = self.get_mc_data()
                if data[0] == UPDATE_EVENT + 1 and data[1] == UPDATE_OP_REQ:
                    return True
        self._failed_update()
        return False

    def _send_pack(self, pack_index):
        if pack_index < self.pack_count:
            cmd = Protocol.pack(UPDATE_EVENT, UPDATE_OP_DATA, self.app_pack[pack_index])
            self.msg.write(cmd)
            return UPDATE_STA_UPLOAD
        else:
            self.msg.write(Protocol.pack(UPDATE_EVENT, UPDATE_OP_END))
            return UPDATE_STA_END

    def _upload_file(self):
        update_status = UPDATE_STA_UPLOAD
        while update_status == UPDATE_STA_UPLOAD:
            if self.wait_mc_data(4000):
                data = self.get_mc_data()
                if data[0] == UPDATE_EVENT + 1 and data[1] == UPDATE_OP_DATA:
                    self.pack_index = (data[2] << 8) + data[3]
                    update_status = self._send_pack(self.pack_index)
            else:
                self.err_count += 1
                if self.err_count > 5:
                    self._failed_update()
                    return False
                update_status = self._send_pack(self.pack_index)
            self.progress(int(self.pack_index / self.pack_count * 100))
        return True

    def _end_update(self):
        count = 0
        while (count < 5):
            if self.wait_mc_data(1000):
                data = self.get_mc_data()
                if data[0] == UPDATE_EVENT + 1 and data[1] == UPDATE_OP_END:
                    return True
            count += 1
        if count == 5:
            self._failed_update()
        return False

    def update(self, file_name, progress):
        self.progress = progress
        self.init_env()
        self.load_updata_file(file_name)
        self.progress(0)
        if not self._wait_update():
            return False
        if not self._upload_file():
            return False
        if not self._end_update():
            return False
        return True
    
